//===----------------------------------------------------------------------===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#ifndef TEST_ALLOCATOR_H
#define TEST_ALLOCATOR_H

#include <cuda/std/cassert>

#include <climits>
#include <cstddef>
#include <cstdlib>
#include <memory>
#include <new>
#include <type_traits>
#include <utility>

#include "test_macros.h"

template <class Alloc>
inline typename std::allocator_traits<Alloc>::size_type alloc_max_size(Alloc const& a)
{
  typedef std::allocator_traits<Alloc> AT;
  return AT::max_size(a);
}

class test_alloc_base
{
protected:
  static int time_to_throw;

public:
  static int throw_after;
  static int count;
  static int alloc_count;
  static int copied;
  static int moved;
  static int converted;

  const static int destructed_value = -1;
  const static int default_value    = 0;
  const static int moved_value      = INT_MAX;

  static void clear()
  {
    assert(count == 0 && "clearing leaking allocator data?");
    count         = 0;
    time_to_throw = 0;
    alloc_count   = 0;
    throw_after   = INT_MAX;
    clear_ctor_counters();
  }

  static void clear_ctor_counters()
  {
    copied    = 0;
    moved     = 0;
    converted = 0;
  }
};

int test_alloc_base::count         = 0;
int test_alloc_base::time_to_throw = 0;
int test_alloc_base::alloc_count   = 0;
int test_alloc_base::throw_after   = INT_MAX;
int test_alloc_base::copied        = 0;
int test_alloc_base::moved         = 0;
int test_alloc_base::converted     = 0;

template <class T>
class test_allocator : public test_alloc_base
{
  int data_; // participates in equality
  int id_; // unique identifier, doesn't participate in equality
  template <class U>
  friend class test_allocator;

public:
  typedef unsigned size_type;
  typedef int difference_type;
  typedef T value_type;
  typedef value_type* pointer;
  typedef const value_type* const_pointer;
  typedef typename std::add_lvalue_reference<value_type>::type reference;
  typedef typename std::add_lvalue_reference<const value_type>::type const_reference;

  template <class U>
  struct rebind
  {
    typedef test_allocator<U> other;
  };

  test_allocator() noexcept
      : data_(0)
      , id_(0)
  {
    ++count;
  }
  explicit test_allocator(int i, int id = 0) noexcept
      : data_(i)
      , id_(id)
  {
    ++count;
  }
  test_allocator(const test_allocator& a) noexcept
      : data_(a.data_)
      , id_(a.id_)
  {
    ++count;
    ++copied;
    assert(a.data_ != destructed_value && a.id_ != destructed_value && "copying from destroyed allocator");
  }
  test_allocator(test_allocator&& a) noexcept
      : data_(a.data_)
      , id_(a.id_)
  {
    ++count;
    ++moved;
    assert(a.data_ != destructed_value && a.id_ != destructed_value && "moving from destroyed allocator");
    a.data_ = moved_value;
    a.id_   = moved_value;
  }
  template <class U>
  test_allocator(const test_allocator<U>& a) noexcept
      : data_(a.data_)
      , id_(a.id_)
  {
    ++count;
    ++converted;
  }
  ~test_allocator() noexcept
  {
    assert(data_ >= 0);
    assert(id_ >= 0);
    --count;
    data_ = destructed_value;
    id_   = destructed_value;
  }
  pointer address(reference x) const
  {
    return &x;
  }
  const_pointer address(const_reference x) const
  {
    return &x;
  }
  pointer allocate(size_type n, const void* = 0)
  {
    assert(data_ >= 0);
    if (time_to_throw >= throw_after)
    {
#if TEST_HAS_EXCEPTIONS()
      throw std::bad_alloc();
#else
      std::terminate();
#endif
    }
    ++time_to_throw;
    ++alloc_count;
    return (pointer)::operator new(n * sizeof(T));
  }
  void deallocate(pointer p, size_type) noexcept
  {
    assert(data_ >= 0);
    --alloc_count;
    ::operator delete((void*) p);
  }
  size_type max_size() const noexcept
  {
    return UINT_MAX / sizeof(T);
  }
  template <class U>
  void construct(pointer p, U&& val)
  {
    ::new (static_cast<void*>(p)) T(std::forward<U>(val));
  }
  void destroy(pointer p) noexcept
  {
    p->~T();
  }
  friend bool operator==(const test_allocator& x, const test_allocator& y)
  {
    return x.data_ == y.data_;
  }
  friend bool operator!=(const test_allocator& x, const test_allocator& y)
  {
    return !(x == y);
  }

  int get_data() const
  {
    return data_;
  }
  int get_id() const
  {
    return id_;
  }
};

template <class T>
class non_default_test_allocator : public test_alloc_base
{
  int data_;

  template <class U>
  friend class non_default_test_allocator;

public:
  typedef unsigned size_type;
  typedef int difference_type;
  typedef T value_type;
  typedef value_type* pointer;
  typedef const value_type* const_pointer;
  typedef typename std::add_lvalue_reference<value_type>::type reference;
  typedef typename std::add_lvalue_reference<const value_type>::type const_reference;

  template <class U>
  struct rebind
  {
    typedef non_default_test_allocator<U> other;
  };

  //    non_default_test_allocator() noexcept : data_(0) {++count;}
  explicit non_default_test_allocator(int i) noexcept
      : data_(i)
  {
    ++count;
  }
  non_default_test_allocator(const non_default_test_allocator& a) noexcept
      : data_(a.data_)
  {
    ++count;
  }
  template <class U>
  non_default_test_allocator(const non_default_test_allocator<U>& a) noexcept
      : data_(a.data_)
  {
    ++count;
  }
  ~non_default_test_allocator() noexcept
  {
    assert(data_ >= 0);
    --count;
    data_ = -1;
  }
  pointer address(reference x) const
  {
    return &x;
  }
  const_pointer address(const_reference x) const
  {
    return &x;
  }
  pointer allocate(size_type n, const void* = 0)
  {
    assert(data_ >= 0);
    if (time_to_throw >= throw_after)
    {
#if TEST_HAS_EXCEPTIONS()
      throw std::bad_alloc();
#else
      std::terminate();
#endif
    }
    ++time_to_throw;
    ++alloc_count;
    return (pointer)::operator new(n * sizeof(T));
  }
  void deallocate(pointer p, size_type) noexcept
  {
    assert(data_ >= 0);
    --alloc_count;
    ::operator delete((void*) p);
  }
  size_type max_size() const noexcept
  {
    return UINT_MAX / sizeof(T);
  }
  template <class U>
  void construct(pointer p, U&& val)
  {
    ::new (static_cast<void*>(p)) T(std::forward<U>(val));
  }
  void destroy(pointer p) noexcept
  {
    p->~T();
  }

  friend bool operator==(const non_default_test_allocator& x, const non_default_test_allocator& y)
  {
    return x.data_ == y.data_;
  }
  friend bool operator!=(const non_default_test_allocator& x, const non_default_test_allocator& y)
  {
    return !(x == y);
  }
};

template <>
class test_allocator<void> : public test_alloc_base
{
  int data_;
  int id_;

  template <class U>
  friend class test_allocator;

public:
  typedef unsigned size_type;
  typedef int difference_type;
  typedef void value_type;
  typedef value_type* pointer;
  typedef const value_type* const_pointer;

  template <class U>
  struct rebind
  {
    typedef test_allocator<U> other;
  };

  test_allocator() noexcept
      : data_(0)
      , id_(0)
  {}
  explicit test_allocator(int i, int id = 0) noexcept
      : data_(i)
      , id_(id)
  {}
  test_allocator(const test_allocator& a) noexcept
      : data_(a.data_)
      , id_(a.id_)
  {}
  template <class U>
  test_allocator(const test_allocator<U>& a) noexcept
      : data_(a.data_)
      , id_(a.id_)
  {}
  ~test_allocator() noexcept
  {
    data_ = -1;
    id_   = -1;
  }

  int get_id() const
  {
    return id_;
  }
  int get_data() const
  {
    return data_;
  }

  friend bool operator==(const test_allocator& x, const test_allocator& y)
  {
    return x.data_ == y.data_;
  }
  friend bool operator!=(const test_allocator& x, const test_allocator& y)
  {
    return !(x == y);
  }
};

template <class T>
class other_allocator
{
  int data_;

  template <class U>
  friend class other_allocator;

public:
  typedef T value_type;

  other_allocator()
      : data_(-1)
  {}
  explicit other_allocator(int i)
      : data_(i)
  {}
  template <class U>
  other_allocator(const other_allocator<U>& a)
      : data_(a.data_)
  {}
  T* allocate(std::size_t n)
  {
    return (T*) ::operator new(n * sizeof(T));
  }
  void deallocate(T* p, std::size_t) noexcept
  {
    ::operator delete((void*) p);
  }

  other_allocator select_on_container_copy_construction() const
  {
    return other_allocator(-2);
  }

  friend bool operator==(const other_allocator& x, const other_allocator& y)
  {
    return x.data_ == y.data_;
  }
  friend bool operator!=(const other_allocator& x, const other_allocator& y)
  {
    return !(x == y);
  }

  typedef std::true_type propagate_on_container_copy_assignment;
  typedef std::true_type propagate_on_container_move_assignment;
  typedef std::true_type propagate_on_container_swap;
};

struct Ctor_Tag
{};

template <typename T>
class TaggingAllocator;

struct Tag_X
{
  // All constructors must be passed the Tag type.

  // DefaultInsertable into vector<X, TaggingAllocator<X>>,
  Tag_X(Ctor_Tag) {}
  // CopyInsertable into vector<X, TaggingAllocator<X>>,
  Tag_X(Ctor_Tag, const Tag_X&) {}
  // MoveInsertable into vector<X, TaggingAllocator<X>>, and
  Tag_X(Ctor_Tag, Tag_X&&) {}

  // EmplaceConstructible into vector<X, TaggingAllocator<X>> from args.
  template <typename... Args>
  Tag_X(Ctor_Tag, Args&&...)
  {}

  // not DefaultConstructible, CopyConstructible or MoveConstructible.
  Tag_X()             = delete;
  Tag_X(const Tag_X&) = delete;
  Tag_X(Tag_X&&)      = delete;

  // CopyAssignable.
  Tag_X& operator=(const Tag_X&)
  {
    return *this;
  }

  // MoveAssignable.
  Tag_X& operator=(Tag_X&&)
  {
    return *this;
  }

private:
  // Not Destructible.
  ~Tag_X() {}

  // Erasable from vector<X, TaggingAllocator<X>>.
  friend class TaggingAllocator<Tag_X>;
};

template <typename T>
class TaggingAllocator
{
public:
  using value_type   = T;
  TaggingAllocator() = default;

  template <typename U>
  TaggingAllocator(const TaggingAllocator<U>&)
  {}

  T* allocate(std::size_t n)
  {
    return std::allocator<T>{}.allocate(n);
  }

  void deallocate(T* p, std::size_t n) noexcept
  {
    std::allocator<T>{}.deallocate(p, n);
  }

  template <typename... Args>
  void construct(Tag_X* p, Args&&... args)
  {
    ::new ((void*) p) Tag_X(Ctor_Tag{}, std::forward<Args>(args)...);
  }

  template <typename U, typename... Args>
  void construct(U* p, Args&&... args)
  {
    ::new ((void*) p) U(std::forward<Args>(args)...);
  }

  template <typename U, typename... Args>
  void destroy(U* p) noexcept
  {
    p->~U();
  }
};

template <typename T, typename U>
bool operator==(const TaggingAllocator<T>&, const TaggingAllocator<U>&)
{
  return true;
}

template <typename T, typename U>
bool operator!=(const TaggingAllocator<T>&, const TaggingAllocator<U>&)
{
  return false;
}

template <std::size_t MaxAllocs>
struct limited_alloc_handle
{
  std::size_t outstanding_;
  void* last_alloc_;

  limited_alloc_handle()
      : outstanding_(0)
      , last_alloc_(nullptr)
  {}

  template <class T>
  T* allocate(std::size_t N)
  {
    if (N + outstanding_ > MaxAllocs)
    {
      TEST_THROW(std::bad_alloc());
    }
    last_alloc_ = ::operator new(N * sizeof(T));
    outstanding_ += N;
    return static_cast<T*>(last_alloc_);
  }

  void deallocate(void* ptr, std::size_t N) noexcept
  {
    if (ptr == last_alloc_)
    {
      last_alloc_ = nullptr;
      assert(outstanding_ >= N);
      outstanding_ -= N;
    }
    ::operator delete(ptr);
  }
};

template <class T, std::size_t N>
class limited_allocator
{
  template <class U, std::size_t UN>
  friend class limited_allocator;
  typedef limited_alloc_handle<N> BuffT;
  std::shared_ptr<BuffT> handle_;

public:
  typedef T value_type;
  typedef value_type* pointer;
  typedef const value_type* const_pointer;
  typedef value_type& reference;
  typedef const value_type& const_reference;
  typedef std::size_t size_type;
  typedef std::ptrdiff_t difference_type;

  template <class U>
  struct rebind
  {
    typedef limited_allocator<U, N> other;
  };

  limited_allocator()
      : handle_(new BuffT)
  {}

  limited_allocator(limited_allocator const& other)
      : handle_(other.handle_)
  {}

  template <class U>
  explicit limited_allocator(limited_allocator<U, N> const& other)
      : handle_(other.handle_)
  {}

private:
  limited_allocator& operator=(const limited_allocator&); // = delete;

public:
  pointer allocate(size_type n)
  {
    return handle_->template allocate<T>(n);
  }
  void deallocate(pointer p, size_type n) noexcept
  {
    handle_->deallocate(p, n);
  }
  size_type max_size() const
  {
    return N;
  }

  BuffT* getHandle() const
  {
    return handle_.get();
  }
};

template <class T, class U, std::size_t N>
inline bool operator==(limited_allocator<T, N> const& LHS, limited_allocator<U, N> const& RHS)
{
  return LHS.getHandle() == RHS.getHandle();
}

template <class T, class U, std::size_t N>
inline bool operator!=(limited_allocator<T, N> const& LHS, limited_allocator<U, N> const& RHS)
{
  return !(LHS == RHS);
}

#endif // TEST_ALLOCATOR_H
